/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.extensions.sql.impl.transform;



import com.bff.gaia.unified.sdk.coders.*;

import com.bff.gaia.unified.sdk.extensions.sql.impl.transform.agg.CovarianceFn;

import com.bff.gaia.unified.sdk.extensions.sql.impl.transform.agg.VarianceFn;

import com.bff.gaia.unified.sdk.extensions.sql.impl.utils.CalciteUtils;

import com.bff.gaia.unified.sdk.schemas.Schema;

import com.bff.gaia.unified.sdk.transforms.*;

import com.bff.gaia.unified.sdk.values.KV;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.ImmutableMap;

import com.bff.gaia.unified.sdk.coders.BigDecimalCoder;

import com.bff.gaia.unified.sdk.coders.BigEndianIntegerCoder;

import com.bff.gaia.unified.sdk.coders.Coder;

import com.bff.gaia.unified.sdk.coders.CoderRegistry;

import com.bff.gaia.unified.sdk.coders.KvCoder;

import com.bff.gaia.unified.sdk.transforms.Combine;

import com.bff.gaia.unified.sdk.transforms.Count;

import com.bff.gaia.unified.sdk.transforms.Max;

import com.bff.gaia.unified.sdk.transforms.Min;

import com.bff.gaia.unified.sdk.transforms.Sum;



import javax.annotation.Nullable;

import java.math.BigDecimal;

import java.math.MathContext;

import java.math.RoundingMode;

import java.util.Map;

import java.util.function.Function;



/** Built-in aggregations functions for COUNT/MAX/MIN/SUM/AVG/VAR_POP/VAR_SAMP. */

public class UnifiedBuiltinAggregations {



  public static final Map<String, Function<Schema.FieldType, Combine.CombineFn<?, ?, ?>>>

      BUILTIN_AGGREGATOR_FACTORIES =

          ImmutableMap.<String, Function<Schema.FieldType, Combine.CombineFn<?, ?, ?>>>builder()

              .put("COUNT", typeName -> Count.combineFn())

              .put("MAX", UnifiedBuiltinAggregations::createMax)

              .put("MIN", UnifiedBuiltinAggregations::createMin)

              .put("SUM", UnifiedBuiltinAggregations::createSum)

              .put("$SUM0", UnifiedBuiltinAggregations::createSum)

              .put("AVG", UnifiedBuiltinAggregations::createAvg)

              .put("VAR_POP", t -> VarianceFn.newPopulation(t.getTypeName()))

              .put("VAR_SAMP", t -> VarianceFn.newSample(t.getTypeName()))

              .put("COVAR_POP", t -> CovarianceFn.newPopulation(t.getTypeName()))

              .put("COVAR_SAMP", t -> CovarianceFn.newSample(t.getTypeName()))

              .build();



  private static MathContext mc = new MathContext(10, RoundingMode.HALF_UP);



  public static Combine.CombineFn<?, ?, ?> create(String functionName, Schema.FieldType fieldType) {



    Function<Schema.FieldType, Combine.CombineFn<?, ?, ?>> aggregatorFactory =

        BUILTIN_AGGREGATOR_FACTORIES.get(functionName);



    if (aggregatorFactory != null) {

      return aggregatorFactory.apply(fieldType);

    }



    throw new UnsupportedOperationException(

        String.format("Aggregator [%s] is not supported", functionName));

  }



  /** {@link Combine.CombineFn} for MAX based on {@link Max} and {@link Combine.BinaryCombineFn}. */

  static Combine.CombineFn createMax(Schema.FieldType fieldType) {

    if (CalciteUtils.isDateTimeType(fieldType)) {

      return new CustMax<>();

    }

    switch (fieldType.getTypeName()) {

      case BOOLEAN:

      case INT16:

      case BYTE:

      case FLOAT:

      case DATETIME:

      case DECIMAL:

      case STRING:

        return new CustMax<>();

      case INT32:

        return Max.ofIntegers();

      case INT64:

        return Max.ofLongs();

      case DOUBLE:

        return Max.ofDoubles();

      default:

        throw new UnsupportedOperationException(

            String.format("[%s] is not support in MAX", fieldType));

    }

  }



  /** {@link Combine.CombineFn} for MIN based on {@link Min} and {@link Combine.BinaryCombineFn}. */

  static Combine.CombineFn createMin(Schema.FieldType fieldType) {

    if (CalciteUtils.isDateTimeType(fieldType)) {

      return new CustMin();

    }

    switch (fieldType.getTypeName()) {

      case BOOLEAN:

      case BYTE:

      case INT16:

      case FLOAT:

      case DATETIME:

      case DECIMAL:

      case STRING:

        return new CustMin();

      case INT32:

        return Min.ofIntegers();

      case INT64:

        return Min.ofLongs();

      case DOUBLE:

        return Min.ofDoubles();

      default:

        throw new UnsupportedOperationException(

            String.format("[%s] is not support in MIN", fieldType));

    }

  }



  /** {@link Combine.CombineFn} for Sum based on {@link Sum} and {@link Combine.BinaryCombineFn}. */

  static Combine.CombineFn createSum(Schema.FieldType fieldType) {

    switch (fieldType.getTypeName()) {

      case INT32:

        return Sum.ofIntegers();

      case INT16:

        return new ShortSum();

      case BYTE:

        return new ByteSum();

      case INT64:

        return Sum.ofLongs();

      case FLOAT:

        return new FloatSum();

      case DOUBLE:

        return Sum.ofDoubles();

      case DECIMAL:

        return new BigDecimalSum();

      default:

        throw new UnsupportedOperationException(

            String.format("[%s] is not support in SUM", fieldType));

    }

  }



  /** {@link Combine.CombineFn} for AVG. */

  static Combine.CombineFn createAvg(Schema.FieldType fieldType) {

    switch (fieldType.getTypeName()) {

      case INT32:

        return new IntegerAvg();

      case INT16:

        return new ShortAvg();

      case BYTE:

        return new ByteAvg();

      case INT64:

        return new LongAvg();

      case FLOAT:

        return new FloatAvg();

      case DOUBLE:

        return new DoubleAvg();

      case DECIMAL:

        return new BigDecimalAvg();

      default:

        throw new UnsupportedOperationException(

            String.format("[%s] is not support in AVG", fieldType));

    }

  }



  static class CustMax<T extends Comparable<T>> extends Combine.BinaryCombineFn<T> {

    @Override

    public T apply(T left, T right) {

      return (right == null || right.compareTo(left) < 0) ? left : right;

    }

  }



  static class CustMin<T extends Comparable<T>> extends Combine.BinaryCombineFn<T> {

    @Override

    public T apply(T left, T right) {

      return (left == null || left.compareTo(right) < 0) ? left : right;

    }

  }



  static class ShortSum extends Combine.BinaryCombineFn<Short> {

    @Override

    public Short apply(Short left, Short right) {

      return (short) (left + right);

    }

  }



  static class ByteSum extends Combine.BinaryCombineFn<Byte> {

    @Override

    public Byte apply(Byte left, Byte right) {

      return (byte) (left + right);

    }

  }



  static class FloatSum extends Combine.BinaryCombineFn<Float> {

    @Override

    public Float apply(Float left, Float right) {

      return left + right;

    }

  }



  static class BigDecimalSum extends Combine.BinaryCombineFn<BigDecimal> {

    @Override

    public BigDecimal apply(BigDecimal left, BigDecimal right) {

      return left.add(right);

    }

  }



  /** {@link Combine.CombineFn} for <em>AVG</em> on {@link Number} types. */

  abstract static class Avg<T extends Number> extends Combine.CombineFn<T, KV<Integer, BigDecimal>, T> {

    @Override

    public KV<Integer, BigDecimal> createAccumulator() {

      return KV.of(0, BigDecimal.ZERO);

    }



    @Override

    public KV<Integer, BigDecimal> addInput(KV<Integer, BigDecimal> accumulator, T input) {

      return KV.of(accumulator.getKey() + 1, accumulator.getValue().add(toBigDecimal(input)));

    }



    @Override

    public KV<Integer, BigDecimal> mergeAccumulators(

        Iterable<KV<Integer, BigDecimal>> accumulators) {

      int size = 0;

      BigDecimal acc = BigDecimal.ZERO;

      for (KV<Integer, BigDecimal> ele : accumulators) {

        size += ele.getKey();

        acc = acc.add(ele.getValue());

      }

      return KV.of(size, acc);

    }



    @Override

    public Coder<KV<Integer, BigDecimal>> getAccumulatorCoder(

		CoderRegistry registry, Coder<T> inputCoder) {

      return KvCoder.of(BigEndianIntegerCoder.of(), BigDecimalCoder.of());

    }



    protected BigDecimal prepareOutput(KV<Integer, BigDecimal> accumulator) {

      return accumulator.getValue().divide(new BigDecimal(accumulator.getKey()), mc);

    }



    @Override

    public abstract T extractOutput(KV<Integer, BigDecimal> accumulator);



    public abstract BigDecimal toBigDecimal(T record);

  }



  static class IntegerAvg extends Avg<Integer> {

    @Override

    @Nullable

    public Integer extractOutput(KV<Integer, BigDecimal> accumulator) {

      return accumulator.getKey() == 0 ? null : prepareOutput(accumulator).intValue();

    }



    @Override

    public BigDecimal toBigDecimal(Integer record) {

      return new BigDecimal(record);

    }

  }



  static class LongAvg extends Avg<Long> {

    @Override

    @Nullable

    public Long extractOutput(KV<Integer, BigDecimal> accumulator) {

      return accumulator.getKey() == 0 ? null : prepareOutput(accumulator).longValue();

    }



    @Override

    public BigDecimal toBigDecimal(Long record) {

      return new BigDecimal(record);

    }

  }



  static class ShortAvg extends Avg<Short> {

    @Override

    @Nullable

    public Short extractOutput(KV<Integer, BigDecimal> accumulator) {

      return accumulator.getKey() == 0 ? null : prepareOutput(accumulator).shortValue();

    }



    @Override

    public BigDecimal toBigDecimal(Short record) {

      return new BigDecimal(record);

    }

  }



  static class ByteAvg extends Avg<Byte> {

    @Override

    @Nullable

    public Byte extractOutput(KV<Integer, BigDecimal> accumulator) {

      return accumulator.getKey() == 0 ? null : prepareOutput(accumulator).byteValue();

    }



    @Override

    public BigDecimal toBigDecimal(Byte record) {

      return new BigDecimal(record);

    }

  }



  static class FloatAvg extends Avg<Float> {

    @Override

    @Nullable

    public Float extractOutput(KV<Integer, BigDecimal> accumulator) {

      return accumulator.getKey() == 0 ? null : prepareOutput(accumulator).floatValue();

    }



    @Override

    public BigDecimal toBigDecimal(Float record) {

      return new BigDecimal(record);

    }

  }



  static class DoubleAvg extends Avg<Double> {

    @Override

    @Nullable

    public Double extractOutput(KV<Integer, BigDecimal> accumulator) {

      return accumulator.getKey() == 0 ? null : prepareOutput(accumulator).doubleValue();

    }



    @Override

    public BigDecimal toBigDecimal(Double record) {

      return new BigDecimal(record);

    }

  }



  static class BigDecimalAvg extends Avg<BigDecimal> {

    @Override

    @Nullable

    public BigDecimal extractOutput(KV<Integer, BigDecimal> accumulator) {

      return accumulator.getKey() == 0 ? null : prepareOutput(accumulator);

    }



    @Override

    public BigDecimal toBigDecimal(BigDecimal record) {

      return record;

    }

  }

}